# semantic_splitter_demo.py
import os
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_openai import ChatOpenAI
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np
from typing import List, Dict, Any, Tuple
import json
import re


class SemanticAwareSplitter:
    """语义感知分割器"""

    def __init__(self):
        # 初始化嵌入模型
        self.embeddings = DashScopeEmbeddings(
            model="text-embedding-v1",
            dashscope_api_key=os.getenv("DASHSCOPE_API_KEY")
        )

        # 初始化LLM
        self.llm = ChatOpenAI(
            api_key=os.getenv("DASHSCOPE_API_KEY"),
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
            model="qwen-turbo"
        )

        # 语义阈值
        self.semantic_threshold = 0.7

    def split_by_semantic_similarity(self, text: str, chunk_size: int = 500,
                                     overlap_ratio: float = 0.1) -> List[Dict[str, Any]]:
        """基于语义相似度进行分割"""

        try:
            # 1. 预分割：按段落分割
            paragraphs = self._extract_paragraphs(text)

            if len(paragraphs) <= 1:
                return [{"content": text, "semantic_score": 1.0, "chunk_type": "single"}]

            # 2. 计算段落嵌入
            paragraph_embeddings = self._get_paragraph_embeddings(paragraphs)

            # 3. 计算语义相似度矩阵
            similarity_matrix = self._calculate_similarity_matrix(paragraph_embeddings)

            # 4. 基于相似度进行智能分组
            semantic_chunks = self._group_by_semantic_similarity(
                paragraphs, similarity_matrix, chunk_size
            )

            # 5. 添加语义重叠
            final_chunks = self._add_semantic_overlap(semantic_chunks, overlap_ratio)

            return final_chunks

        except Exception as e:
            print(f"语义分割失败: {e}")
            # 回退到基础分割
            fallback_splitter = RecursiveCharacterTextSplitter(
                chunk_size=chunk_size,
                chunk_overlap=int(chunk_size * overlap_ratio)
            )
            basic_chunks = fallback_splitter.split_text(text)
            return [{"content": chunk, "semantic_score": 0.5, "chunk_type": "fallback"}
                    for chunk in basic_chunks]

    def _extract_paragraphs(self, text: str) -> List[str]:
        """提取段落"""
        # 按双换行分割段落
        paragraphs = text.split('\n\n')

        # 清理并过滤空段落
        cleaned_paragraphs = []
        for para in paragraphs:
            cleaned = para.strip()
            if cleaned and len(cleaned) > 20:  # 过滤太短的段落
                cleaned_paragraphs.append(cleaned)

        # 如果段落太少，尝试按单换行分割
        if len(cleaned_paragraphs) < 3:
            lines = text.split('\n')
            current_para = ""

            for line in lines:
                line = line.strip()
                if line:
                    current_para += line + " "
                    # 如果当前段落够长，就分割
                    if len(current_para) > 100 and line.endswith(('。', '.', '！', '?')):
                        cleaned_paragraphs.append(current_para.strip())
                        current_para = ""

            # 添加最后一个段落
            if current_para.strip():
                cleaned_paragraphs.append(current_para.strip())

        return cleaned_paragraphs

    def _get_paragraph_embeddings(self, paragraphs: List[str]) -> List[List[float]]:
        """获取段落嵌入向量"""
        try:
            # 批量计算嵌入
            embeddings = self.embeddings.embed_documents(paragraphs)
            return embeddings
        except Exception as e:
            print(f"嵌入计算失败: {e}")
            # 返回随机向量作为备选
            return [[0.0] * 1024 for _ in paragraphs]

    def _calculate_similarity_matrix(self, embeddings: List[List[float]]) -> np.ndarray:
        """计算相似度矩阵"""
        embeddings_array = np.array(embeddings)
        similarity_matrix = cosine_similarity(embeddings_array)
        return similarity_matrix

    def _group_by_semantic_similarity(self, paragraphs: List[str],
                                      similarity_matrix: np.ndarray,
                                      target_chunk_size: int) -> List[Dict[str, Any]]:
        """基于语义相似度分组"""

        chunks = []
        used_paragraphs = set()

        for i, para in enumerate(paragraphs):
            if i in used_paragraphs:
                continue

            # 开始新的语义块
            current_chunk = [para]
            current_size = len(para)
            used_paragraphs.add(i)

            # 查找相似的段落
            similarities = [(j, similarity_matrix[i][j])
                            for j in range(len(paragraphs))
                            if j not in used_paragraphs]

            # 按相似度排序
            similarities.sort(key=lambda x: x[1], reverse=True)

            # 添加相似段落直到达到目标大小
            for j, sim_score in similarities:
                if sim_score < self.semantic_threshold:
                    break

                para_j = paragraphs[j]
                if current_size + len(para_j) <= target_chunk_size * 1.2:  # 允许20%的弹性
                    current_chunk.append(para_j)
                    current_size += len(para_j)
                    used_paragraphs.add(j)
                else:
                    break

            # 创建语义块
            chunk_content = '\n\n'.join(current_chunk)
            avg_similarity = np.mean([similarity_matrix[i][j] for j in used_paragraphs if j != i]) if len(
                used_paragraphs) > 1 else 1.0

            chunks.append({
                "content": chunk_content,
                "semantic_score": float(avg_similarity),
                "paragraph_count": len(current_chunk),
                "chunk_type": "semantic"
            })

        return chunks

    def _add_semantic_overlap(self, chunks: List[Dict[str, Any]],
                              overlap_ratio: float) -> List[Dict[str, Any]]:
        """添加语义重叠"""

        if len(chunks) <= 1 or overlap_ratio <= 0:
            return chunks

        enhanced_chunks = []

        for i, chunk in enumerate(chunks):
            content = chunk["content"]

            # 添加前一个块的结尾作为上下文
            if i > 0:
                prev_content = chunks[i - 1]["content"]
                overlap_size = int(len(prev_content) * overlap_ratio)
                if overlap_size > 0:
                    prev_overlap = prev_content[-overlap_size:]
                    content = prev_overlap + "\n\n" + content

            # 添加后一个块的开头作为上下文
            if i < len(chunks) - 1:
                next_content = chunks[i + 1]["content"]
                overlap_size = int(len(next_content) * overlap_ratio)
                if overlap_size > 0:
                    next_overlap = next_content[:overlap_size]
                    content = content + "\n\n" + next_overlap

            enhanced_chunk = chunk.copy()
            enhanced_chunk["content"] = content
            enhanced_chunk["has_overlap"] = i > 0 or i < len(chunks) - 1

            enhanced_chunks.append(enhanced_chunk)

        return enhanced_chunks


def demonstrate_semantic_splitting():
    """演示语义分割"""

    print("🧠 语义分割技术演示")
    print("=" * 25)

    # 复杂简历示例
    complex_resume = """
李工程师个人简历

基本信息
我叫李工程师，今年30岁，拥有8年的软件开发经验。目前居住在北京，联系电话是13800138000，邮箱是li@email.com。我专注于后端开发和系统架构设计，特别在分布式系统和微服务架构方面有深入研究。

教育背景和学术成就
我2014年毕业于清华大学计算机科学与技术专业，获得工学学士学位。在校期间，我主修了数据结构、算法设计、操作系统、计算机网络等核心课程，GPA达到3.8/4.0。我的毕业设计是关于分布式系统一致性算法的研究，获得了优秀毕业论文奖。在学习期间，我还参与了多个开源项目，积累了丰富的编程经验。

早期职业经历和成长
毕业后，我加入了百度，担任初级软件工程师职位。在百度的两年时间里，我主要负责搜索引擎后端服务的开发和维护。我从最基础的代码编写开始，逐步学会了大规模分布式系统的设计和运维。期间我参与了搜索系统的性能优化项目，使得查询响应时间降低了30%。这段经历让我深刻理解了互联网大厂的技术体系和工程文化。

技术能力和专业技能
在编程语言方面，我精通Java和Python，有8年的实际开发经验。Java方面，我熟悉Spring全家桶，包括Spring Boot、Spring Cloud、Spring Security等框架，能够快速构建微服务架构的应用。Python方面，我擅长使用Django和Flask进行Web开发，也有丰富的数据处理和机器学习相关经验。

在数据库技术方面，我熟练使用MySQL进行关系型数据建模和查询优化，有丰富的索引设计和性能调优经验。同时我也掌握Redis、MongoDB等NoSQL数据库的使用，能够根据业务场景选择合适的存储方案。在大数据处理方面，我有Hadoop、Spark、Kafka等技术的实践经验。

架构设计和系统优化能力
在系统架构方面，我有丰富的微服务架构设计经验。我参与设计和实现了多个大型分布式系统，包括电商平台、金融支付系统、内容推荐系统等。我深入理解CAP理论、分布式一致性、服务治理、容错处理等核心概念，能够根据业务需求设计高可用、高性能的系统架构。

在性能优化方面，我有丰富的经验。我曾经将一个电商系统的吞吐量从每秒1000请求提升到10000请求，主要通过缓存策略优化、数据库查询优化、异步处理、负载均衡等手段实现。我也有丰富的JVM调优经验，能够分析和解决内存泄漏、GC性能等问题。

领导力和团队管理经验
在最近的三年中，我担任了技术团队负责人的角色，管理一个15人的开发团队。我负责技术决策、项目规划、人员培养等工作。我建立了完善的代码审查机制、技术分享制度、新人培训体系，显著提升了团队的技术水平和工作效率。

我也有丰富的跨部门协作经验，经常需要与产品、设计、测试、运维等不同角色的同事协作。我善于沟通，能够准确理解业务需求，并将其转化为可执行的技术方案。我也经常参与技术选型和架构评审，为公司的技术发展方向提供建议。

创新项目和技术探索
我一直保持对新技术的关注和学习。最近两年，我深入研究了云原生技术，包括Docker、Kubernetes、Service Mesh等。我主导了公司的容器化改造项目，将传统的单体应用迁移到Kubernetes平台，显著提升了系统的可扩展性和运维效率。

我也在人工智能领域有一些探索。我利用业余时间学习了机器学习和深度学习技术，并在实际项目中应用了推荐算法、自然语言处理等技术。我相信AI技术将会深刻改变软件开发的方式，我希望能够在这个领域继续深入发展。

职业规划和发展目标
对于未来的职业发展，我希望能够在技术和管理两个方向都有所成长。在技术方面，我计划深入学习云计算、人工智能、区块链等前沿技术，成为技术专家。在管理方面，我希望能够带领更大的团队，参与更复杂的项目，最终成为优秀的技术管理者。

我也希望能够为开源社区做出更多贡献，分享自己的技术经验和思考。我相信技术的发展需要整个行业的共同努力，我愿意成为这个过程中的积极参与者。
"""

    # 初始化语义分割器
    semantic_splitter = SemanticAwareSplitter()

    # 1. 语义分割演示
    print("1️⃣ 语义分割处理：")

    semantic_chunks = semantic_splitter.split_by_semantic_similarity(
        complex_resume,
        chunk_size=800,
        overlap_ratio=0.15
    )

    print(f"   原文长度: {len(complex_resume)}字符")
    print(f"   语义分割结果: {len(semantic_chunks)}个块")

    for i, chunk in enumerate(semantic_chunks, 1):
        content = chunk["content"]
        semantic_score = chunk["semantic_score"]
        chunk_type = chunk["chunk_type"]

        print(f"\n   块{i} ({chunk_type}):")
        print(f"   长度: {len(content)}字符")
        print(f"   语义分数: {semantic_score:.3f}")

        # 识别内容主题
        topic = identify_content_topic(content)
        print(f"   主题: {topic}")

        # 显示内容预览
        preview = content.replace('\n', ' ').strip()
        print(f"   内容: {preview}...")

    # 2. 与传统分割对比
    print("\n2️⃣ 语义分割 vs 传统分割对比：")

    # 传统递归分割
    traditional_splitter = RecursiveCharacterTextSplitter(
        chunk_size=800,
        chunk_overlap=120,
        separators=["\n\n", "\n", "。", "；", " "]
    )

    traditional_chunks = traditional_splitter.split_text(complex_resume)

    print(f"   传统分割: {len(traditional_chunks)}个块")
    print(f"   语义分割: {len(semantic_chunks)}个块")

    # 比较分割质量
    traditional_quality = evaluate_splitting_quality(traditional_chunks, "traditional")
    semantic_quality = evaluate_splitting_quality(
        [chunk["content"] for chunk in semantic_chunks], "semantic"
    )

    print(f"\n   质量对比:")
    print(f"   传统分割 - 语义完整性: {traditional_quality['semantic_completeness']:.2f}")
    print(f"   语义分割 - 语义完整性: {semantic_quality['semantic_completeness']:.2f}")
    print(f"   传统分割 - 长度一致性: {traditional_quality['length_consistency']:.2f}")
    print(f"   语义分割 - 长度一致性: {semantic_quality['length_consistency']:.2f}")


def identify_content_topic(content: str) -> str:
    """识别内容主题"""

    content_lower = content.lower()

    # 主题关键词映射
    topic_keywords = {
        "基本信息": ["姓名", "年龄", "联系", "电话", "邮箱", "住址"],
        "教育背景": ["教育", "学历", "毕业", "大学", "专业", "学位", "GPA"],
        "工作经历": ["工作", "职位", "公司", "经历", "担任", "就职"],
        "技术技能": ["技术", "编程", "语言", "框架", "数据库", "开发"],
        "项目经验": ["项目", "系统", "开发", "设计", "实现", "架构"],
        "管理经验": ["团队", "管理", "领导", "负责人", "协作", "培养"],
        "职业规划": ["未来", "发展", "目标", "计划", "希望", "愿景"]
    }

    # 计算每个主题的匹配分数
    topic_scores = {}
    for topic, keywords in topic_keywords.items():
        score = sum(1 for keyword in keywords if keyword in content_lower)
        if score > 0:
            topic_scores[topic] = score

    # 返回得分最高的主题
    if topic_scores:
        return max(topic_scores.items(), key=lambda x: x[1])[0]
    else:
        return "其他内容"


def evaluate_splitting_quality(chunks: List[str], split_type: str) -> Dict[str, float]:
    """评估分割质量"""

    if not chunks:
        return {"semantic_completeness": 0, "length_consistency": 0}

    # 1. 语义完整性评估
    semantic_scores = []
    for chunk in chunks:
        # 检查是否以完整的语义单位结束
        chunk_stripped = chunk.strip()

        # 句子完整性
        sentence_complete = chunk_stripped.endswith(('。', '！', '？', '.', '!', '?'))

        # 段落完整性
        paragraph_complete = '\n\n' in chunk or len(chunk.split('\n')) <= 3

        # 主题一致性（简化检查）
        topic_keywords = identify_content_topic(chunk)
        topic_consistent = topic_keywords != "其他内容"

        semantic_score = (
                sentence_complete * 0.4 +
                paragraph_complete * 0.3 +
                topic_consistent * 0.3
        )
        semantic_scores.append(semantic_score)

    # 2. 长度一致性评估
    lengths = [len(chunk) for chunk in chunks]
    avg_length = sum(lengths) / len(lengths)
    length_variance = sum((l - avg_length) ** 2 for l in lengths) / len(lengths)
    length_consistency = max(0, 1 - length_variance / (avg_length ** 2))

    return {
        "semantic_completeness": sum(semantic_scores) / len(semantic_scores),
        "length_consistency": length_consistency
    }


# 运行演示
demonstrate_semantic_splitting()
